package com.colin.util.util;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import java.util.Date;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.*;

public class CustThreadPoolExecutor extends ThreadPoolExecutor {

    private static final Logger logger = LoggerFactory.getLogger("CustThreadPoolExecutor");

    // 保存任务开始执行的时间,当任务结束时,用任务结束时间减去开始时间计算任务执行时间
    private ConcurrentHashMap<String, Date> startTimes = new ConcurrentHashMap<>(128);

    public CustThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime,
                                  TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory,
                                  RejectedExecutionHandler rejectedExecutionHandler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, rejectedExecutionHandler);
    }

    /**
     * 所有线程都会委托给这个execute方法，在这个方法中我们把父线程的MDC内容赋值给子线程
     */
    @Override
    public void execute(Runnable runnable) {
        // 获取父线程MDC中的内容，必须在run方法之前，否则等异步线程执行的时候有可能MDC里面的值已经被清空了，这个时候就会返回null
        Map<String, String> context = MDC.getCopyOfContextMap();
        super.execute(() -> run(runnable, context));
    }


    @Override
    public Future<?> submit(Runnable task) {
        // 获取父线程MDC中的内容，必须在run方法之前，否则等异步线程执行的时候有可能MDC里面的值已经被清空了，这个时候就会返回null
        Map<String, String> context = MDC.getCopyOfContextMap();
        return super.submit(() -> run(task, context));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        // 获取父线程MDC中的内容，必须在run方法之前，否则等异步线程执行的时候有可能MDC里面的值已经被清空了，这个时候就会返回null
        Map<String, String> context = MDC.getCopyOfContextMap();
        return super.submit(() -> run(task, context));
    }

    @Override
    protected void beforeExecute(Thread t, Runnable r) {
        startTimes.put(String.valueOf(r.hashCode()), new Date());
        logger.info("子线程开始执行");
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
        Date startDate = startTimes.remove(String.valueOf(r.hashCode()));
        Date finishDate = new Date();
        long diff = finishDate.getTime();
        if (Objects.nonNull(startDate)) {
            diff -= startDate.getTime();
        }
        logger.info("当前线程池监控:Duration: {} ms,PoolSize: {}, CorePoolSize: {}, Active: {}, "
                        + "Completed: {}, Task: {}, Queue: {}, LargestPoolSize: {}, MaximumPoolSize: {},"
                        + "KeepAliveTime: {}, isShutdown: {}, isTerminated: {}",
                diff, this.getPoolSize(), this.getCorePoolSize(), this.getActiveCount(),
                this.getCompletedTaskCount(), this.getTaskCount(), this.getQueue().size(),
                this.getLargestPoolSize(), this.getMaximumPoolSize(),
                this.getKeepAliveTime(TimeUnit.MILLISECONDS), this.isShutdown(), this.isTerminated(), t);
    }

    /**
     * 子线程委托的执行方法
     *
     * @param callable {@link Runnable}
     * @param context  父线程MDC内容
     */
    private <T> T run(Callable<T> callable, Map<String, String> context) throws Exception {
        // 将父线程的MDC内容传给子线程
        if (context != null) {
            try {
                MDC.setContextMap(context);
            } catch (Exception e) {
                logger.error(e.getMessage(), e);
            }
        }
        T call;
        try {
            // 执行异步操作
            call = callable.call();
        } finally {
            // 清空MDC内容
            MDC.clear();
        }
        return call;
    }

    /**
     * 子线程委托的执行方法
     *
     * @param runnable {@link Runnable}
     * @param context  父线程MDC内容
     */
    private void run(Runnable runnable, Map<String, String> context) {
        // 将父线程的MDC内容传给子线程
        if (context != null) {
            try {
                MDC.setContextMap(context);
            } catch (Exception e) {
                logger.error(e.getMessage(), e);
            }
        }
        try {
            // 执行异步操作
            runnable.run();
        } finally {
            // 清空MDC内容
            MDC.clear();
        }
    }
}
